import pandas as pd
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import matplotlib.font_manager as fm
import os

from _00下载数据 import get_today_date

def to_str(dt):
    return f"{dt.year}年{dt.month}月"

def quarter_mean(se):
    """
    计算季度的平均值
    """
    month = se.index[-1].month
    year = se.index[-1].year
    if month in [1,2,3]:
        quarter = 1
    elif month in [4,5,6]:
        quarter = 2
    elif month in [7,8,9]:
        quarter = 3
    else:
        quarter = 4
    ss = f"{year}年Q{quarter}均值"
    se.index = [to_str(dt) for dt in se.index]
    se[ss] = se.mean()
    return se


def split_series_by_quarters(se):
    """
    将原始序列分割成本季度、上个季度和上上个季度的三个序列
    
    参数:
    se: pandas Series对象，索引为日期类型
    
    返回:
    tuple: (本季度序列, 上个季度序列, 上上个季度序列)
    """
    # 确保索引是日期类型
    if not pd.api.types.is_datetime64_any_dtype(se.index):
        se.index = pd.to_datetime(se.index)
    
    # 获取当前日期和季度
    today = datetime.today()
    current_year = today.year
    current_quarter = (today.month - 1) // 3 + 1
    
    # 计算本季度的开始和结束日期
    current_quarter_start = datetime(current_year, 3 * (current_quarter - 1) + 1, 1)
    if current_quarter == 4:
        current_quarter_end = datetime(current_year + 1, 1, 1) - timedelta(days=1)
    else:
        current_quarter_end = datetime(current_year, 3 * current_quarter + 1, 1) - timedelta(days=1)
    
    # 计算上个季度的开始和结束日期
    if current_quarter == 1:
        prev_quarter_start = datetime(current_year - 1, 10, 1)
        prev_quarter_end = datetime(current_year, 1, 1) - timedelta(days=1)
    else:
        prev_quarter_start = datetime(current_year, 3 * (current_quarter - 2) + 1, 1)
        prev_quarter_end = datetime(current_year, 3 * (current_quarter - 1) + 1, 1) - timedelta(days=1)
    
    # 计算上上个季度的开始和结束日期
    if current_quarter == 1:
        prev_prev_quarter_start = datetime(current_year - 1, 7, 1)
        prev_prev_quarter_end = datetime(current_year - 1, 10, 1) - timedelta(days=1)
    elif current_quarter == 2:
        prev_prev_quarter_start = datetime(current_year - 1, 10, 1)
        prev_prev_quarter_end = datetime(current_year, 1, 1) - timedelta(days=1)
    else:
        prev_prev_quarter_start = datetime(current_year, 3 * (current_quarter - 3) + 1, 1)
        prev_prev_quarter_end = datetime(current_year, 3 * (current_quarter - 2) + 1, 1) - timedelta(days=1)
    
    # 根据日期范围筛选序列
    current_quarter_se = se[(se.index >= current_quarter_start) & (se.index <= current_quarter_end)]
    prev_quarter_se = se[(se.index >= prev_quarter_start) & (se.index <= prev_quarter_end)]
    prev_prev_quarter_se = se[(se.index >= prev_prev_quarter_start) & (se.index <= prev_prev_quarter_end)]
    current_quarter_se = quarter_mean(current_quarter_se)
    prev_quarter_se = quarter_mean(prev_quarter_se)
    prev_prev_quarter_se = quarter_mean(prev_prev_quarter_se)
    return current_quarter_se, prev_quarter_se, prev_prev_quarter_se


def plot_quarterly_comparison(prev_prev_q, prev_q, current_q, title='季度经济指标比较', save_path=None):
    """
    绘制季度比较柱状图，用竖线区分不同季度
    
    参数:
    current_q: 本季度序列
    prev_q: 上个季度序列
    prev_prev_q: 上上个季度序列
    title: 图表标题
    save_path: 图表保存路径，默认为None（不保存）
    """
    # 创建画布和子图
    plt.figure(figsize=(12, 6))
    
    # 绘制不同季度的数据点，使用不同颜色
    
    if not prev_prev_q.empty:
        # 检查并处理均值标签
        for idx, value in prev_prev_q.items():
            if '均值' in idx:
                plt.bar(idx, value, color='#ff69b4', alpha=0.9, label='上上个季度均值' if '均值' in idx else '上上个季度')
            else:
                plt.bar(idx, value, color='#2ca02c', alpha=0.9, label='上上个季度')
    
    if not prev_q.empty:
        # 检查并处理均值标签
        for idx, value in prev_q.items():
            if '均值' in idx:
                plt.bar(idx, value, color='#ff69b4', alpha=0.9, label='上个季度均值' if '均值' in idx else '上个季度')
            else:
                plt.bar(idx, value, color='#ff7f0e', alpha=0.9, label='上个季度')
    
    if not current_q.empty:
        # 检查并处理均值标签
        for idx, value in current_q.items():
            if '均值' in idx:
                plt.bar(idx, value, color='#ff69b4', alpha=0.9, label='本季度均值' if '均值' in idx else '本季度')
                plt.text(idx, value, f'{value:.2f}', ha='center', va='bottom', fontsize=9)
            else:
                plt.bar(idx, value, color='#1f77b4', alpha=0.9, label='本季度')
    
    # 添加标签和标题
    plt.title(title, fontsize=14)
    plt.xlabel('日期', fontsize=12)
    plt.ylabel('指标值', fontsize=12)
    # plt.legend(fontsize=10)
    # 绘制红色水平虚线，y=2.5的位置
    plt.axhline(y=current_q[-1], color='r', linestyle='--')
    # 设置日期格式
    plt.xticks(rotation=45)
    
    # 调整布局
    plt.tight_layout()
    
    # 保存图表（如果提供了保存路径）
    if save_path:
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        plt.savefig(save_path, dpi=300)
        print(f"图表已保存至: {save_path}")
    
    # 显示图表
    plt.show()

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'WenQuanYi Micro Hei', 'Heiti TC']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


def plot_three_quarter(pth, name):    
    today = get_today_date()
    # pth = '生产测_工业'
    # name = '工业增加值当月同比（合并1-2月）'
    path = f'{pth}_{today}.xlsx'

    df = pd.read_excel(f"cache/{path}", index_col=0)
    se = df[name]
    current_q, prev_q, prev_prev_q = split_series_by_quarters(se)
    print("本季度数据:")
    print(current_q)
    print("\n上个季度数据:")
    print(prev_q)
    print("\n上上个季度数据:")
    print(prev_prev_q)

    # 生成并显示季度比较柱状图
    plot_title = f'{name}季度比较图'

    # 创建保存路径
    plot_dir = f'plot/plot_{today}'
    os.makedirs(plot_dir, exist_ok=True)
    save_path = f'{plot_dir}/{name}_季度比较图.png'

    # 绘制图表
    plot_quarterly_comparison(
        prev_prev_q, 
        prev_q, 
        current_q, 
        title=plot_title,
        )


# 使用示例
if __name__ == '__main__':
    # pth = '生产测_工业'
    # name = '工业增加值当月同比（合并1-2月）'
    # plot_three_quarter(pth, name)
    # pth = '生产测_服务业'
    # name = '服务业生产指数：当月同比'
    # plot_three_quarter(pth, name)
    plot_three_quarter('需求侧_消费', '社会消费品零售总额：当月同比')